'''
Created on Nov 26, 2011

@author: oabalbin
'''

import sys
import subprocess
import numpy as np
from optparse import OptionParser
from collections import deque, defaultdict

                
def read_positions(ifile):
    '''
    '''
    ifile = open(ifile)
    for l in ifile:
        if l.startswith('#'):
            continue
        fields=l.strip('\n').split(',')
        fields.pop()
    
    ifile.close()
    return fields


def isec_arrays(filea,fileb):
    '''
    '''
    m1=set(read_positions(filea))
    m2=set(read_positions(fileb))
    #print sys.getsizeof(m1), sys.getsizeof(m2)
    isec=m1.intersection(m2)        
    return isec

def isec_snvs_annot_union(snvs_union, snvs_annot_dbs):
    '''
    This functions obtains all query snvs that are 
    common to a particular database, ex: dbSNP
    '''
    pm = defaultdict()
    for db, snv_db in snvs_annot_dbs.iteritems():
        pm[db] = set(snvs_union).intersection(set(snv_db))
           
    return pm

def union_snvs_dict(union_snvs):
    '''
    It gives an index to each snp. 
    So it can be used later to find arrays
    '''
    pm = defaultdict()
    for i, snv in enumerate(union_snvs):
        pm[snv]=i
    return pm


if __name__ == '__main__':
    '''
    Write to an array the positions in a particualr vcf.
    
    '''
    optionparser = OptionParser("usage: %prog [options] ")
    optionparser.add_option("-f", "--pos_files", action="append", type="str",
                            dest="pos_files",
                            help="others_vcf_files. Pass vcf files without .gz extention even if the file .gz is there")

    optionparser.add_option("-o", "--out_file", type="str",dest="ofile",
                            help="others_vcf_files. Pass vcf files without .gz extention even if the file .gz is there")
    
    optionparser.add_option("-b", "--base", dest="base", action="store_true", default=False,
                            help="determine if the ouput file is position or pos.base")
    optionparser.add_option("-g", "--geno", dest="geno", action="store_true", default=False,
                            help="determine if the ouput file is position or pos.base.geno")

    
    
    (options, args) = optionparser.parse_args()
    
    
    
    snvs_union = set()
    
    pfiles=defaultdict()
    snvs_annot_dbs = defaultdict()

    # Read pos.base.geno files
    First=True
    for ft in options.pos_files:
        # Generate a positions files for query vcf.
        print ft
        thsnvs = read_positions(ft)
        snvs_annot_dbs[ft]=thsnvs
        snvs_union.update(set(thsnvs))
        # isec set
        if First:
            isec_set=set(thsnvs)
            First=False
        else:
            isec_set.intersection_update( set(thsnvs) )
    
    diff_snvs = snvs_union.difference(isec_set)
    '''
    chr14@39272493|C>A|0/1
    '''
    
    # For the snvs shared across all individuals, isec set
    nrows=len(isec_set)+len(diff_snvs)
    ncols=1
    mytype=[('CHRM',int),('POS',int),('IBD',int),('MUT','S3'),('GENO','S3')]
    mytype_ibd=[('CHRM',int),('POS',int),('IBD',int)]
  
    #location_arr=np.empty(( nrows, ncols),dtype=mytype)
    all_snvs=[]
    all_ibd=[]
    ibd=str(1)
    i=0
    for snv in isec_set:
        loc,base,geno = snv.split('|')[0],snv.split('|')[1].replace('>','/'),snv.split('|')[2]
        chr,pos=loc.split('@')[0],loc.split('@')[1]
        chrN=chr.replace('chr','')
        if chrN=='X':
            chrN=23
        elif chrN=='Y':
            chrN=24
        elif chrN=='M':
            chrN=25
             
        all_snvs.append((chrN,pos,ibd,base,geno)) # This indixing is important
        all_ibd.append( (chrN,pos,ibd))
        #print ol, snv
        #print location_arr[i,:]
        i+=1
    print "The length of isec",len( isec_set.intersection(diff_snvs) )
    ibd=str(0)
    for snv in diff_snvs:
        loc,base,geno = snv.split('|')[0],snv.split('|')[1].replace('>','/'),snv.split('|')[2]
        chr,pos=loc.split('@')[0],loc.split('@')[1]        
        chrN=chr.replace('chr','')

        if chrN=='X':
            chrN=23
        elif chrN=='Y':
            chrN=24
        elif chrN=='M':
            chrN=25
        all_snvs.append((chrN,pos,ibd,base,geno)) 
        all_ibd.append( (chrN,pos,ibd))              
        i+=1
    # Sort the snv table
    location_arr=np.array( all_snvs,dtype=mytype )
    location_arr_sorted = np.sort(location_arr, order=['CHRM', 'POS'])
    
    # This takes time. Look for a better way to do it.    
    ibd_arr=np.array( all_ibd,dtype=mytype_ibd )
    ibd_arr_sorted=np.unique(ibd_arr)
    ibd_arr_sorted=np.sort(ibd_arr_sorted, order=['CHRM', 'POS'])
    
    of=open(options.ofile,'w')   
    of2=open(options.ofile+'_all','w')
    
    # Write all the snps.
    for i in range(len(location_arr_sorted)):
        
        ol=list(location_arr_sorted[i])
        if ol[0]==23:
            ol[0]='X'
        elif ol[0]==24:
            ol[0]='Y'
        elif ol[0]==25: # CHRM M
            continue
        #print ",".join(map(str,ol)).replace(',','\t')
        of2.write('chr'+",".join(map(str,ol)).replace(',','\t')+'\n')
    of2.close()           

    # Write only the unique snps when ibd is considered
    for i in range(len(ibd_arr_sorted)):
        
        ol=list(ibd_arr_sorted[i])
        if ol[0]==23:
            ol[0]='X'
        elif ol[0]==24:
            ol[0]='Y'
        elif ol[0]==25: # CHRM M
            continue
        #print ",".join(map(str,ol)).replace(',','\t')
        of.write('chr'+",".join(map(str,ol)).replace(',','\t')+'\n')

    of.close()           
    
    '''    
    ########
    ibd=str(1)
    for snv in isec_set:
        loc,base,geno = snv.split('|')[0],snv.split('|')[1].replace('>','/'),snv.split('|')[2]
        chr,pos=loc.split('@')[0],loc.split('@')[1]
        chrN=chr.replace('chr','')
        ol=['chr',chrN,pos,ibd,base,geno]
        of.write(",".join(ol).replace(',','\t')+'\n')
    
    # For the snvs not shared across all the individuals
    ibd=str(0)
    for snv in diff_snvs:
        loc,base,geno = snv.split('|')[0],snv.split('|')[1].replace('>','/'),snv.split('|')[2]
        chr,pos=loc.split('@')[0],loc.split('@')[1]
        chrN=chr.replace('chr','')
        ol=['chr',chrN,pos,ibd,base,geno]
        of.write(",".join(ol).replace(',','\t')+'\n')
    
    of.close()
    
    st=options.ofile+'_sorted'
    f=open(st,'w')
    rc=subprocess.call(['sort',options.ofile],stdout=f)
    f.close()
    
    if rc==0:
        ifile=open(st)
        of=open(st+'_ibd_ready','w')
        for l in ifile:
            f=l.strip('\n').split('\t')
            chr,chrN,pos,base,geno=f[0],f[1],f[2],f[3],f[4]
            ol=[chr+chrN,pos,base,geno]
            of.write(",".join(ol).replace(',','\t')+'\n')
    '''    

    


    
    
